import numpy as np  
import matplotlib.pyplot as plt  
from matplotlib.widgets import Slider  
  
# 定义7自由度机械臂模型（考虑姿态）  
class SevenDOFRobotModel:  
    def __init__(self, lengths):  
        self.lengths = lengths  # 各连杆长度  
        self.n_joints = len(lengths)  # 关节数量  
      
    # 正向运动学（考虑位置和姿态）  
    def forward_kinematics(self, theta):  
        pos = np.zeros(3)  # 位置 (x, y, z)  
        ori = np.zeros(3)  # 姿态 (roll, pitch, yaw)  
        current_angle = 0  
        for i in range(self.n_joints):  
            current_angle += theta[i]  
            # 计算位置增量（简化，假设在xy平面）  
            dx = self.lengths[i] * np.cos(current_angle)  
            dy = self.lengths[i] * np.sin(current_angle)  
            pos += np.array([dx, dy, 0])  
            # 计算姿态增量（简化，假设绕z轴旋转）  
            ori[2] += theta[i]  
        return np.concatenate([pos, ori])  
      
    # 计算雅可比矩阵  
    def jacobian(self, theta):  
        J = np.zeros((6, self.n_joints))  
        current_angle = 0  
        for i in range(self.n_joints):  
            current_angle += theta[i]  
            # 位置雅可比  
            J[0, i] = -self.lengths[i] * np.sin(current_angle)  
            J[1, i] = self.lengths[i] * np.cos(current_angle)  
            # 姿态雅可比（简化，假设绕z轴旋转）  
            J[5, i] = 1  
        return J  

# 雅可比迭代求解逆运动学（带阻尼和加权）  
def inverse_kinematics_jacobian(target_pose, robot, initial_theta, max_iter=1000, tol=1e-3, damping=0.1, weights=None):  
    theta = np.array(initial_theta, dtype=np.float64)  
    if weights is None:  
        weights = np.ones(robot.n_joints)  # 默认权重为1  
    for i in range(max_iter):  
        current_pose = robot.forward_kinematics(theta)  
        error = target_pose - current_pose  
        if np.linalg.norm(error) < tol:  
            return theta, True  
        J = robot.jacobian(theta)  
        # 计算加权雅可比矩阵的伪逆  
        J_weighted = J * weights  
        #J_pinv = np.linalg.pinv(J_weighted)  
        #J_pinv = J.T @ np.linalg.inv(J @ J.T + damping**2 * np.eye(6))  # 使用SVD求伪逆
        #print(f"J_pinv shape: {J_pinv.shape}")
        J_pinv = np.linalg.inv(J.T @ J + damping**2 * np.eye(7)) @ J.T  # 使用正规方程求伪逆
        delta_theta = np.dot(J_pinv, error)  
        # 应用阻尼系数  
        delta_theta_damped = delta_theta * damping  
        theta += delta_theta_damped  
    return theta, False  
  
# 初始化图形和坐标系  
fig, ax = plt.subplots()  
plt.subplots_adjust(bottom=0.25, left=0.25)  # 为滑块和文本留出空间  
  
# 设置坐标系的范围和标签  
ax.set_xlim(-10, 10)  
ax.set_ylim(-10, 10)  
ax.set_xlabel('X')  
ax.set_ylabel('Y')  
ax.grid(True)  
  
# 初始机械臂参数  
lengths = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]  
robot = SevenDOFRobotModel(lengths)  
  
# 初始关节角度  
initial_theta = [np.pi/4, np.pi/4, np.pi/4, np.pi/4, np.pi/4, np.pi/4, np.pi/4]  
current_theta = np.array(initial_theta, dtype=np.float64)  
  
# 计算初始机械臂端点位置  
arm_end = robot.forward_kinematics(current_theta)[:2]  # 只取位置信息  
  
# 绘制机械臂的函数  
def draw_arm(theta):  
    end_point = robot.forward_kinematics(theta)[:2]  # 只取位置信息  
    x = [0]  
    y = [0]  
    current_x, current_y = 0, 0  
    for i in range(len(theta)):  
        current_x += lengths[i] * np.cos(np.sum(theta[:i+1]))  
        current_y += lengths[i] * np.sin(np.sum(theta[:i+1]))  
        x.append(current_x)  
        y.append(current_y)  
    line.set_data(x, y)  # 更新机械臂端点坐标  
    fig.canvas.draw_idle()  # 重绘图形  
  
# 创建机械臂图形对象  
line, = ax.plot([0, arm_end[0]], [0, arm_end[1]], 'ro-')  
  
# 创建滑块  
ax_slider_x = plt.axes([0.25, 0.1, 0.65, 0.03])  # X轴滑块的位置和大小  
slider_x = Slider(ax_slider_x, 'Target X', -8, 8, valinit=arm_end[0], valstep=0.1)  
  
ax_slider_y = plt.axes([0.25, 0.05, 0.65, 0.03])  # Y轴滑块的位置和大小  
slider_y = Slider(ax_slider_y, 'Target Y', -8, 8, valinit=arm_end[1], valstep=0.1)  
  
# 创建文本对象用于显示当前坐标和关节角度  
text = ax.text(-9, 9, f'Current Position: ({arm_end[0]:.1f}, {arm_end[1]:.1f})\nTheta: {current_theta}', visible=True)  
  
# 滑块事件处理函数  
def update(val):  
    # 获取当前滑块值作为目标位置  
    target_pos = np.array([slider_x.val, slider_y.val, 0])  # 目标位置（z=0）  
    target_ori = np.array([0, 0, 0])  # 目标姿态（简化）  
    target_pose = np.concatenate([target_pos, target_ori])  
      
    # 求解逆运动学  
    global current_theta  
    current_theta, success = inverse_kinematics_jacobian(target_pose, robot, current_theta)  # 使用上一次的结果作为初始猜测  
      
    if not success:
        print("Inverse kinematics failed: Target pose is unreachable.")
        return  # 直接返回，不更新机械臂状态

    # 更新机械臂端点位置  
    draw_arm(current_theta)  
      
    # 更新显示的坐标和关节角度
    end_point_current = robot.forward_kinematics(current_theta)[:2]  # 只取位置信息  
    text.set_text(f'Current Position: ({end_point_current[0]:.1f}, {end_point_current[1]:.1f})\nTheta: {np.round(current_theta, 3)}')  
    fig.canvas.draw_idle()  
  
# 将滑块与更新函数关联  
slider_x.on_changed(update)  
slider_y.on_changed(update)  
  
plt.show()